from __future__ import annotations

import numpy as np
import vtk
from dataclasses import dataclass
from typing import List, Optional

from matplotlib import cm
from PyQt5 import QtWidgets
from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor

from config.settings import settings
from services.registration_manager import RegisteredGroup, RegistrationManager


@dataclass
class GroupState:
    data: RegisteredGroup
    ply_visible: bool = False
    obj_visible: bool = False


class ViewerApp(QtWidgets.QDialog):
    def __init__(
        self,
        groups: Optional[List[RegisteredGroup]] = None,
        parent: Optional[QtWidgets.QWidget] = None,
        registration_manager: Optional[RegistrationManager] = None,
    ) -> None:
        super().__init__(parent)
        self.setWindowTitle("点云误差分析")
        self.setGeometry(200, 200, 1000, 700)

        self.registration_manager = registration_manager or RegistrationManager()
        loaded_groups = groups if groups is not None else self.registration_manager.load_registered_groups()
        self.group_states = [GroupState(data=group) for group in loaded_groups]
        self.actors = {}
        self.show_only_error = False
        self.scalar_bar = None

        self._setup_style()
        self._build_ui()

        self._update_scene()
        self._update_error_info()

        self.show()
        self.interactor.Initialize()

    # ------------------------------------------------------------------ #
    # UI assembly
    # ------------------------------------------------------------------ #
    def _setup_style(self) -> None:
        self.setStyleSheet(
            """
            QDialog, QWidget {
                background-color: #2D2D2D;
                color: #E0E0E0;
            }
            QLabel {
                color: #FFFFFF;
            }
            QPushButton {
                background-color: #505050;
                border: 1px solid #606060;
                border-radius: 8px;
                padding: 10px;
                font-size: 14px;
                color: #FFFFFF;
            }
            QPushButton:hover {
                background-color: #606060;
            }
            QPushButton:pressed {
                background-color: #404040;
            }
            QTextEdit {
                background-color: #232323;
                color: #E0E0E0;
                border: 1px solid #404040;
                border-radius: 4px;
                font-size: 12pt;
            }
        """
        )

    def _build_ui(self) -> None:
        main_layout = QtWidgets.QHBoxLayout(self)
        main_layout.setSpacing(10)
        main_layout.setContentsMargins(10, 10, 10, 10)

        self.vtk_widget = QVTKRenderWindowInteractor(self)
        main_layout.addWidget(self.vtk_widget, stretch=3)

        control_panel_layout = QtWidgets.QVBoxLayout()
        control_panel_layout.setSpacing(10)
        control_panel_layout.setContentsMargins(10, 10, 10, 10)

        self.error_toggle_btn = QtWidgets.QPushButton("仅显示误差点云")
        self.error_toggle_btn.setCheckable(True)
        self.error_toggle_btn.clicked.connect(self._toggle_error_display)
        control_panel_layout.addWidget(self.error_toggle_btn)

        group_label = QtWidgets.QLabel("组选择：")
        group_label.setStyleSheet("font-weight: bold; font-size: 12pt;")
        control_panel_layout.addWidget(group_label)
        self._init_checkboxes(control_panel_layout)

        error_label = QtWidgets.QLabel("误差分析信息：")
        error_label.setStyleSheet("font-weight: bold; font-size: 12pt;")
        control_panel_layout.addWidget(error_label)

        self.error_text = QtWidgets.QTextEdit()
        self.error_text.setReadOnly(True)
        self.error_text.setMinimumHeight(180)
        control_panel_layout.addWidget(self.error_text)

        self.status_label = QtWidgets.QLabel("就绪")
        self.status_label.setStyleSheet("color: #AAAAAA; font-size: 10px;")
        control_panel_layout.addWidget(self.status_label)

        control_panel_layout.addStretch(1)

        right_widget = QtWidgets.QWidget()
        right_widget.setLayout(control_panel_layout)
        main_layout.addWidget(right_widget, stretch=2)

        self.renderer = vtk.vtkRenderer()
        self.vtk_widget.GetRenderWindow().AddRenderer(self.renderer)
        self.interactor = self.vtk_widget.GetRenderWindow().GetInteractor()

    # ------------------------------------------------------------------ #
    # UI interactions
    # ------------------------------------------------------------------ #
    def _toggle_error_display(self) -> None:
        self.show_only_error = self.error_toggle_btn.isChecked()
        self._update_scene()
        if self.scalar_bar is not None:
            self.renderer.AddActor2D(self.scalar_bar)

    def _init_checkboxes(self, layout: QtWidgets.QVBoxLayout) -> None:
        for state in self.group_states:
            group_widget = QtWidgets.QWidget()
            group_h_layout = QtWidgets.QHBoxLayout(group_widget)
            group_h_layout.setSpacing(10)
            group_h_layout.setContentsMargins(0, 0, 0, 0)

            label = QtWidgets.QLabel(state.data.name)
            label.setFixedWidth(100)
            label.setStyleSheet("font-size: 12pt;")
            group_h_layout.addWidget(label)

            ply_checkbox = QtWidgets.QCheckBox("PLY")
            obj_checkbox = QtWidgets.QCheckBox("OBJ")
            ply_checkbox.setStyleSheet("font-size: 12pt;")
            obj_checkbox.setStyleSheet("font-size: 12pt;")

            ply_checkbox.stateChanged.connect(self._make_checkbox_callback(state, "ply"))
            obj_checkbox.stateChanged.connect(self._make_checkbox_callback(state, "obj"))

            group_h_layout.addWidget(ply_checkbox)
            group_h_layout.addWidget(obj_checkbox)
            layout.addWidget(group_widget)

    def _make_checkbox_callback(self, state: GroupState, geo_type: str):
        def callback(check_state: int) -> None:
            if geo_type == "ply":
                state.ply_visible = check_state == 2
            else:
                state.obj_visible = check_state == 2
            self._update_scene()
            self._update_error_info()

        return callback

    # ------------------------------------------------------------------ #
    # Rendering helpers
    # ------------------------------------------------------------------ #
    def _create_scalar_bar(self, min_error: float, max_error: float) -> None:
        if self.scalar_bar is not None:
            self.renderer.RemoveActor2D(self.scalar_bar)

        lut = vtk.vtkLookupTable()
        lut.SetNumberOfTableValues(256)
        lut.SetRange(min_error, max_error)
        lut.Build()

        colormap = cm.get_cmap("coolwarm")
        for i in range(256):
            color = colormap(i / 255.0)[:3]
            lut.SetTableValue(i, color[0], color[1], color[2], 1.0)

        self.scalar_bar = vtk.vtkScalarBarActor()
        self.scalar_bar.SetLookupTable(lut)
        self.scalar_bar.SetTitle("误差 (m)")
        self.scalar_bar.SetNumberOfLabels(2)
        self.scalar_bar.SetLabelFormat("%.3f")
        self.scalar_bar.SetPosition(0.9, 0.1)
        self.scalar_bar.SetWidth(0.08)
        self.scalar_bar.SetHeight(0.8)

        self.renderer.AddActor2D(self.scalar_bar)

    def _update_scene(self) -> None:
        self.renderer.RemoveAllViewProps()
        if self.scalar_bar is not None:
            self.renderer.RemoveActor2D(self.scalar_bar)

        max_error = 0.0
        min_error = float("inf")
        has_error_points = False

        for state in self.group_states:
            if state.ply_visible:
                actor = self._create_point_cloud_actor(state.data)
                if actor is not None:
                    self.renderer.AddActor(actor)
                    self.actors[f"{state.data.name}_ply"] = actor

                distances = state.data.distances
                above = distances > state.data.threshold
                if np.any(above):
                    has_error_points = True
                    max_error = max(max_error, float(np.max(distances[above])))
                    min_error = min(min_error, float(np.min(distances[above])))

            if state.obj_visible:
                actor = self._create_mesh_actor(state.data.mesh)
                self.renderer.AddActor(actor)
                self.actors[f"{state.data.name}_obj"] = actor

        if has_error_points:
            self._create_scalar_bar(min_error, max_error)

        self.renderer.ResetCamera()
        self.vtk_widget.GetRenderWindow().Render()

    def _create_point_cloud_actor(self, group: RegisteredGroup) -> Optional[vtk.vtkActor]:
        ply = group.ply
        if not ply.has_points():
            return None

        points = vtk.vtkPoints()
        vertices = vtk.vtkCellArray()
        colors = vtk.vtkUnsignedCharArray()
        colors.SetNumberOfComponents(3)
        colors.SetName("Colors")

        distances = group.distances
        for idx, point in enumerate(ply.points):
            if self.show_only_error and distances[idx] <= group.threshold:
                continue

            pid = points.InsertNextPoint(point[0], point[1], point[2])
            vertices.InsertNextCell(1)
            vertices.InsertCellPoint(pid)

            color = ply.colors[idx] if ply.has_colors() else [0.5, 0.5, 0.5]
            colors.InsertNextTuple3(int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))

        polydata = vtk.vtkPolyData()
        polydata.SetPoints(points)
        polydata.SetVerts(vertices)
        polydata.GetPointData().SetScalars(colors)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(polydata)

        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        actor.GetProperty().SetPointSize(settings.viewer.point_size)
        return actor

    def _create_mesh_actor(self, mesh) -> vtk.vtkActor:
        points = vtk.vtkPoints()
        for pt in mesh.vertices:
            points.InsertNextPoint(pt[0], pt[1], pt[2])

        triangles = vtk.vtkCellArray()
        for tri in mesh.triangles:
            triangle = vtk.vtkTriangle()
            for i in range(3):
                triangle.GetPointIds().SetId(i, tri[i])
            triangles.InsertNextCell(triangle)

        polydata = vtk.vtkPolyData()
        polydata.SetPoints(points)
        polydata.SetPolys(triangles)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(polydata)

        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        return actor

    def _update_error_info(self) -> None:
        all_distances = []
        total_above_count = 0
        min_error = float("inf")
        max_error = 0.0

        for state in self.group_states:
            if state.ply_visible:
                distances = state.data.distances
                threshold = state.data.threshold
                above = distances > threshold
                if np.any(above):
                    all_distances.extend(distances[above])
                    total_above_count += int(np.sum(above))
                    min_error = min(min_error, float(np.min(distances[above])))
                    max_error = max(max_error, float(np.max(distances[above])))

        if not all_distances:
            self.error_text.setText("")
            return

        all_distances = np.array(all_distances)
        avg_error = float(np.mean(all_distances))

        threshold = self.group_states[0].data.threshold if self.group_states else settings.registration.default_threshold
        error_info = (
            f"平均误差: {avg_error:.4f} m\n"
            f"最大误差: {max_error:.4f} m\n"
            f"最小误差: {min_error:.4f} m\n"
            f"超出阈值点数: {total_above_count}\n"
            f"阈值: {threshold:.4f} m"
        )
        self.error_text.setText(error_info)


def load_groups_with_manager(manager: Optional[RegistrationManager] = None) -> List[RegisteredGroup]:
    """
    Convenience wrapper for loading registered groups via the manager.
    """
    registration_manager = manager or RegistrationManager()
    return registration_manager.load_registered_groups()
